#pragma once

#include <nanobdd/Common.h>
#include <nanobdd/Node.h>
#include <map>
#include <shared_mutex>
#include <unordered_map>
#include <vector>

namespace nanobdd {

class Bucket {
 public:
  Bucket(size_t size) : size_(size), nodes_(size), currSize_(0) {
    mutex_ = std::make_unique<std::shared_mutex>();
    for (int i = 0; i < 34; i++) {
      levelNodeMap_.push_back({});
    }
  }

  // Node*
  // operator()(int level, Node* low, Node* high) {
  //   // const std::lock_guard<std::mutex> lock(*g_i_mutex);
  //   // search in nodes
  //   {
  //     std::shared_lock<std::shared_mutex> lock(*mutex_);
  //     for (int i = 0; i < currSize_; ++i) {
  //       auto& node = nodes_.at(i);
  //       if (node.low() == low && node.high() == high && node.level() ==
  //       level) {
  //         return &node;
  //       }
  //     }
  //   }
  //   std::lock_guard<std::shared_mutex> lock(*mutex_);
  //   // create the node
  //   auto& node = nodes_.at(currSize_);
  //   node.level_ = level;
  //   node.low_ = low;
  //   node.high_ = high;
  //   currSize_++;

  //   return &node;
  // };

  Node*
  operator()(int level, Node* low, Node* high) {
    auto& map = levelNodeMap_.at(level);
    {
      std::shared_lock<std::shared_mutex> lock(*mutex_);
      auto it1 = map.find(low);
      if (it1 != map.end()) {
        auto it2 = it1->second.find(high);
        if (it2 != it1->second.end()) {
          return &(it2->second);
        }
      }
    }
    std::lock_guard<std::shared_mutex> lock(*mutex_);
    auto it1 = map.find(low);
    if (it1 != map.end()) {
      auto it2 = it1->second.find(high);
      if (it2 != it1->second.end()) {
        return &(it2->second);
      }
    }
    auto& node = levelNodeMap_[level][low][high];
    node.level_ = level;
    node.low_ = low;
    node.high_ = high;

    return &node;
  }

 private:
  size_t size_;
  size_t currSize_; // current occupied number of nodes
  std::vector<Node> nodes_;
  std::unique_ptr<std::shared_mutex> mutex_;
  std::map<int, std::map<Node*, std::map<Node*, Node>>> mapNodes_;
  std::vector<std::unordered_map<Node*, std::unordered_map<Node*, Node>>>
      levelNodeMap_;
};

class NodeTable {
 public:
  NodeTable(size_t tableSize, int bucketSize) : tableSize_(tableSize) {
    for (int i = 0; i < tableSize; ++i) {
      // buckets_.push_back(std::move(Bucket(bucketSize)));
      buckets_.emplace_back(bucketSize);
    }
  };

  /**
   * Get or create a node
   */
  Node*
  operator()(int level, Node* l, Node* r) {
    auto hash = TRIPLEpo(level, l, r) % tableSize_;
    auto& bucket = buckets_.at(hash);

    return bucket(level, l, r);
  };

 private:
  size_t tableSize_;
  std::vector<Bucket> buckets_;
};

} // namespace nanobdd